import os
import time
from pathlib import Path
import logging
from openai import OpenAI
from volcenginesdkarkruntime import Ark
import requests
import json
from dotenv import load_dotenv  # 导入 load_dotenv 函数

# 加载 .env 文件
load_dotenv()

# 确保 logs 文件夹存在
logs_dir = Path('logs')
logs_dir.mkdir(parents=True, exist_ok=True)

# 配置日志记录，同时输出到控制台和文件
log_filename = time.strftime("%Y%m%d.log")
log_file_path = logs_dir / log_filename
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')

# 配置文件处理器
file_handler = logging.FileHandler(log_file_path, encoding='utf-8')
file_handler.setFormatter(formatter)

# 配置控制台处理器
console_handler = logging.StreamHandler()
console_handler.setFormatter(formatter)

# 配置日志记录器
logging.basicConfig(level=logging.INFO, handlers=[file_handler, console_handler])

class DoubaoAPI:
    def __init__(self, api_key):
        """初始化豆包 API

        Args:
            api_key (str): 豆包 API key
        """
        self.client = Ark(
            base_url="https://ark.cn-beijing.volces.com/api/v3",
            api_key=api_key
        )
        # 统一模型名称
        self.model = "doubao-pro-256k-241115"

    def get_answer(self, question, prompt=""):
        """调用豆包 API 获取回答

        Args:
            question (str): 问题内容
            prompt (str): 提示词

        Returns:
            str: API 返回的回答
        """
        try:
            # 组合提示词和问题
            full_question = question

            messages = [
                {"role": "system", "content": prompt if prompt else "你是人工智能助手."},
                {"role": "user", "content": full_question}
            ]
            # 记录请求路径和请求体到日志
            base_url = "https://ark.cn-beijing.volces.com/api/v3"
            logging.info(f"[DoubaoAPI] 请求路径: {base_url}/chat/completions")
            logging.info(f"[DoubaoAPI] 请求体: {{'model': '{self.model}', 'messages': {messages}}}")

            completion = self.client.chat.completions.create(
                model=self.model,
                messages=messages
            )
            logging.info(f"[DoubaoAPI] 成功获取回答，长度: {len(completion.choices[0].message.content)}")
            return completion.choices[0].message.content
        except Exception as e:
            logging.error(f"[DoubaoAPI] API 调用出错: {e}")
            return None


class GPTAPI:
    def __init__(self, api_key):
        """初始化 OpenAI API

        Args:
            api_key (str): OpenAI API key
        """
        self.client = OpenAI(api_key=api_key)
        self.model = "gpt-4o"

    def get_answer(self, question, prompt=""):
        """调用 GPT-4o API 获取回答

        Args:
            question (str): 问题内容
            prompt (str): 提示词

        Returns:
            str: API 返回的回答
        """
        try:
            # 组合提示词和问题
            full_question = question

            messages = [
                {"role": "system", "content": prompt if prompt else "你是人工智能助手."},
                {"role": "user", "content": full_question}
            ]

            # 记录请求路径和请求体到日志
            logging.info(f"[GPTAPI] 请求路径: {self.client.base_url}/chat/completions")
            logging.info(f"[GPTAPI] 请求体: {{'model': '{self.model}', 'messages': {messages}, 'temperature': 0.7}}")

            # 修复代码部分
            completion = self.client.chat.completions.create(
                model=self.model,
                messages=messages
            )
            contents = completion.choices[0].message.content
            # 检查 contents 是否为 None
            if contents is not None:
                logging.info(f"[GPTAPI] 成功获取回答，长度: {len(contents)}")
            else:
                logging.info("[GPTAPI] 未获取到有效回答")
            return contents
        except Exception as e:
            logging.error(f"[GPTAPI] API 调用出错: {e}")
            return None


class OpenRouterAPI:
    def __init__(self, api_key, site_url="", site_name="", model="anthropic/claude-3.5-sonnet"):
        """初始化 OpenRouter API

        Args:
            api_key (str): OpenRouter API key
            site_url (str): 网站 URL，用于 openrouter.ai 排名
            site_name (str): 网站名称，用于 openrouter.ai 排名
            model (str): 使用的模型，默认为 "anthropic/claude-3.5-sonnet"
        """
        self.api_key = api_key
        self.site_url = site_url
        self.site_name = site_name
        self.model = model

    def get_answer(self, question, prompt=""):
        """调用 OpenRouter API 获取回答

        Args:
            question (str): 问题内容
            prompt (str): 提示词

        Returns:
            str: API 返回的回答
        """
        try:
            # 组合提示词和问题
            full_question = question

            messages = [
                {"role": "system", "content": prompt if prompt else "你是人工智能助手."},
                {"role": "user", "content": full_question}
            ]

            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
            }
            if self.site_url:
                headers["HTTP-Referer"] = self.site_url
            if self.site_name:
                headers["X-Title"] = self.site_name

            data = {
                "model": self.model,
                "messages": messages
            }

            # 记录请求路径和请求体到日志
            logging.info(f"[OpenRouterAPI] 请求路径: https://openrouter.ai/api/v1/chat/completions")
            logging.info(f"[OpenRouterAPI] 请求体: {json.dumps(data)}")

            response = requests.post(
                url="https://openrouter.ai/api/v1/chat/completions",
                headers=headers,
                data=json.dumps(data)
            )
            response.raise_for_status()
            response_data = response.json()
            # 输出完整的 API 返回结果用于调试
            logging.info(f"[OpenRouterAPI] 完整 API 返回结果: {response_data}")
            answer = response_data["choices"][0]["message"]["content"]
            logging.info(f"[OpenRouterAPI] 成功获取回答，长度: {len(answer)}")
            return answer
        except KeyError as e:
            logging.error(f"[OpenRouterAPI] 解析 API 返回结果时出错，缺少 '{e}' 字段，完整结果: {response.json()}")
            return None
        except Exception as e:
            logging.error(f"[OpenRouterAPI] API 调用出错: {e}")
            return None


class SiliconflowAPI:
    def __init__(self, api_key, model="Qwen/QwQ-32B"):
        """初始化 Siliconflow API

        Args:
            api_key (str): Siliconflow API key
            model (str): 使用的模型，默认为 "Qwen/QwQ-32B"
        """
        self.api_key = api_key
        self.model = model
        self.url = "https://api.siliconflow.cn/v1/chat/completions"

    def get_answer(self, question, prompt=""):
        """调用 Siliconflow API 获取回答

        Args:
            question (str): 问题内容
            prompt (str): 提示词

        Returns:
            str: API 返回的回答
        """
        try:
            messages = [
                {"role": "system", "content": prompt if prompt else "你是人工智能助手."},
                {"role": "user", "content": question}
            ]

            payload = {
                "model": self.model,
                "messages": messages,
                "stream": False,
                "max_tokens": 512,
                "stop": None,
                "temperature": 0.7,
                "top_p": 0.7,
                "top_k": 50,
                "frequency_penalty": 0.5,
                "n": 1,
                "response_format": {"type": "text"},
                "tools": [
                    {
                        "type": "function",
                        "function": {
                            "description": "<string>",
                            "name": "<string>",
                            "parameters": {},
                            "strict": False
                        }
                    }
                ]
            }

            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            }

            # 记录请求路径和请求体到日志
            logging.info(f"[SiliconflowAPI] 请求路径: {self.url}")
            logging.info(f"[SiliconflowAPI] 请求体: {json.dumps(payload)}")

            response = requests.post(self.url, json=payload, headers=headers)
            response.raise_for_status()
            response_data = response.json()
            answer = response_data["choices"][0]["message"]["content"]
            logging.info(f"[SiliconflowAPI] 成功获取回答，长度: {len(answer)}")
            return answer
        except KeyError as e:
            logging.error(f"[SiliconflowAPI] 解析 API 返回结果时出错，缺少 '{e}' 字段，完整结果: {response.json()}")
            return None
        except Exception as e:
            logging.error(f"[SiliconflowAPI] API 调用出错: {e}")
            return None

def process_md_files(source_dir, target_dir, model_type="gpt-4o", **api_keys):
    """处理 MD 文件并获取回答

    Args:
        source_dir (str): 源文件目录
        target_dir (str): 目标文件目录
        model_type (str): 模型类型，可选 "gpt-4o" 或 "doubao" 或 "openrouter" 或 "deepseek-r1"
        **api_keys: API 密钥，根据 model_type 提供相应的密钥
    """
    # 根据模型类型创建 API 实例
    if model_type == "gpt-4o":
        api = GPTAPI(api_keys.get("openai_key"))
    elif model_type == "doubao":
        api = DoubaoAPI(api_keys.get("doubao_key"))
    elif model_type == "openrouter":
        api = OpenRouterAPI(
            api_key=api_keys.get("openrouter_key"),
            site_url=api_keys.get("openrouter_site_url", ""),
            site_name=api_keys.get("openrouter_site_name", ""),
            model=api_keys.get("openrouter_model", "anthropic/claude-3.5-sonnet")
        )
    elif model_type == "deepseek-r1":
        api = SiliconflowAPI(
            api_key=api_keys.get("siliconflow_key"),
            model="deepseek-ai/DeepSeek-R1"
        )
    else:
        logging.error(f"不支持的模型类型: {model_type}")
        return

    # 确保目标目录存在
    Path(target_dir).mkdir(parents=True, exist_ok=True)

    # 读取提示词
    prompt = ""
    prompt_file = Path("chat_prompt.md")
    if prompt_file.exists():
        try:
            with open(prompt_file, 'r', encoding='utf-8') as f:
                prompt = f.read().strip()
            if prompt:
                logging.info(f"已加载提示词，长度：{len(prompt)} 字符")
        except Exception as e:
            logging.error(f"读取提示词文件出错: {e}")

    # 获取所有 .md 文件
    md_files = list(Path(source_dir).glob("*.md"))
    total_files = len(md_files)

    logging.info(f"\n开始处理，共发现 {total_files} 个 MD 文件")

    for index, md_file in enumerate(md_files, 1):
        try:
            logging.info(f"\n[{index}/{total_files}] 处理文件: {md_file.name}")

            # 读取问题内容
            with open(md_file, 'r', encoding='utf-8') as f:
                question = f.read().strip()

            if not question:
                logging.info("-> 跳过: 文件内容为空")
                continue

            # 构建目标文件路径
            target_file = Path(target_dir) / md_file.name

            # 如果目标文件已存在，跳过处理
            if target_file.exists():
                logging.info(f"-> 跳过: 答案文件已存在 ({target_file.name})")
                continue

            logging.info("-> 调用 API 获取回答...")
            answer = api.get_answer(question, prompt)

            if answer:
                # 保存回答
                with open(target_file, 'w', encoding='utf-8') as f:
                    f.write(f"{answer}")
                logging.info("   回答已保存")
            else:
                logging.warning("!! API 返回为空，跳过保存")

            # 添加延时避免触发 API 限制
            time.sleep(1)

        except Exception as e:
            logging.error(f"!! 处理文件出错: {md_file.name}")
            logging.error(f"!! 错误信息: {e}")

    logging.info(f"\n任务完成! 共处理 {total_files} 个文件")


if __name__ == "__main__":
    # 从 .env 文件中获取源目录和目标目录
    source_dir = os.getenv("SOURCE_DIR")
    target_dir = os.getenv("TARGET_DIR")

    if not source_dir or not target_dir:
        logging.error("错误：未在 .env 文件中配置源目录或目标目录。")
        exit(1)

    # 定义模型选项
    model_options = {
        1: "gpt-4o",
        2: "doubao",
        3: "openrouter",
        4: "deepseek-r1"
    }

    # 显示模型选项
    logging.info("请选择要使用的模型：")
    for key, value in model_options.items():
        logging.info(f"{key}. {value}")

    # 获取用户输入
    try:
        choice = int(input("请输入对应的数字: "))
        model_type = model_options.get(choice)
        if not model_type:
            logging.error("错误：输入的数字无效，请输入有效的模型编号。")
            exit(1)
    except ValueError:
        logging.error("错误：输入无效，请输入一个数字。")
        exit(1)

    if model_type == "gpt-4o":
        # 从 .env 文件中获取 OpenAI API key
        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key:
            api_key = input("请输入 OpenAI API Key: ").strip()
        if not api_key:
            logging.error("错误：未提供 OpenAI API Key")
            exit(1)
        process_md_files(source_dir, target_dir,
                         model_type, openai_key=api_key)
    elif model_type == "doubao":
        # 从 .env 文件中获取豆包 API key
        api_key = os.getenv("ARK_API_KEY")
        if not api_key:
            api_key = input("请输入豆包 API Key: ").strip()
        if not api_key:
            logging.error("错误：未提供豆包 API Key")
            exit(1)
        process_md_files(source_dir, target_dir,
                         model_type, doubao_key=api_key)
    elif model_type == "openrouter":
        # 从 .env 文件中获取 OpenRouter API key
        api_key = os.getenv("OPENROUTER_API_KEY")
        if not api_key:
            api_key = input("请输入 OpenRouter API Key: ").strip()
        if not api_key:
            logging.error("错误：未提供 OpenRouter API Key")
            exit(1)
        site_url = os.getenv("OPENROUTER_SITE_URL", "")
        site_name = os.getenv("OPENROUTER_SITE_NAME", "")
        process_md_files(source_dir, target_dir,
                         model_type,
                         openrouter_key=api_key,
                         openrouter_site_url=site_url,
                         openrouter_site_name=site_name)
    elif model_type == "deepseek-r1":
        # 从 .env 文件中获取 Siliconflow API key
        api_key = os.getenv("SILICONFLOW_API_KEY")
        if not api_key:
            api_key = input("请输入 Siliconflow API Key: ").strip()
        if not api_key:
            logging.error("错误：未提供 Siliconflow API Key")
            exit(1)
        process_md_files(source_dir, target_dir,
                         model_type,
                         siliconflow_key=api_key)
